package MLDA_Process;

import gnu.trove.map.hash.TIntIntHashMap;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.apache.commons.math3.distribution.BetaDistribution;
import org.apache.commons.math3.distribution.GammaDistribution;
import org.apache.commons.math3.distribution.NormalDistribution;
import org.apache.commons.math3.random.JDKRandomGenerator;
import org.apache.commons.math3.special.Gamma;

import Commons.IOUtil;

public class SparsePAMForShortTexts_V3 implements Runnable {

    public int K1 = 1000;
    public int K2 = 100;

    public int M;
    public int V;

    public double alpha = 0.1;

    public double alpha0 = 1E-7;
    public double alpha1 = 0.1;

    public double beta = 0.01;

    public double gamma0 = 1.0;
    public double gamma1 = 1.0;

    public boolean a[][];
    public int a_sum[];

    public double pi_a[];

    public JDKRandomGenerator rand;

    public int mp[];

    public int npk[][];
    public int npkSum[];

    public int nkw[][];
    public int nkwSum[];

    public int zAssigns_1[];
    public int zAssigns_2[][];

    public int niters = 3000;


    public int saveStep = 1000;
    public String inputPath = "";
    public String outputPath = "";

    public List<List<Integer>> docs = new ArrayList<List<Integer>>();
    public HashMap<String, Integer> w2i = new HashMap<String, Integer>();
    public HashMap<Integer, String> i2w = new HashMap<Integer, String>();


    public SparsePAMForShortTexts_V3(int P, int K, int iter, int saveStep, double alpha, double alpha0, double alpha1, double beta, String inputPath, String outputPath) {
        this.K1 = P;
        this.K2 = K;
        this.niters = iter;
        this.saveStep = saveStep;
        this.alpha = alpha;
        this.alpha0 = alpha0;
        this.alpha1 = alpha1;
        this.beta = beta;
        this.inputPath = inputPath;
        this.outputPath = outputPath;
    }


    public void loadTxts(String txtPath) {
        BufferedReader reader = IOUtil.getReader(txtPath, "UTF-8");

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();

                String[] tokens = line.trim().split("\\s+");
                for (String token : tokens) {
                    if (!w2i.containsKey(token)) {
                        w2i.put(token, w2i.size());
                        i2w.put(w2i.get(token), token);
                    }
                    doc.add(w2i.get(token));
                }
                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        M = docs.size();
        V = w2i.size();
        return;
    }

    public void initSTPAM() {
        rand = new JDKRandomGenerator();
        rand.setSeed(System.currentTimeMillis());

        BetaDistribution betaDist = new BetaDistribution(rand, gamma1, gamma0);

        mp = new int[K1];

        npk = new int[K1][K2];
        npkSum = new int[K1];

        nkw = new int[K2][V];
        nkwSum = new int[K2];

        a = new boolean[K1][K2];
        a_sum = new int[K1];

        pi_a = new double[K1];

        for (int k1 = 0; k1 != K1; k1++) {
            pi_a[k1] = betaDist.sample();
            for (int k2 = 0; k2 != K2; k2++) {
                a[k1][k2] = true;
            }
            a_sum[k1] = K2;
        }

        zAssigns_1 = new int[M];
        zAssigns_2 = new int[M][];

        for (int m = 0; m != M; m++) {
            int N = docs.get(m).size();
            zAssigns_2[m] = new int[N];

            int z1 = (int) Math.floor(rand.nextDouble() * K1);
            zAssigns_1[m] = z1;

            mp[z1]++;

            for (int n = 0; n != N; n++) {
                int w = docs.get(m).get(n);
                int z2 = (int) Math.floor(rand.nextDouble() * K2);

                npk[z1][z2]++;
                npkSum[z1]++;

                nkw[z2][w]++;
                nkwSum[z2]++;

                zAssigns_2[m][n] = z2;
            }
        }
    }

    public void sampleZ1(int m) {
        int z1 = zAssigns_1[m];
        int N = docs.get(m).size();

        mp[z1]--;

        TIntIntHashMap k2Count = new TIntIntHashMap();
        for (int n = 0; n != N; n++) {
            int z2 = zAssigns_2[m][n];
            k2Count.adjustOrPutValue(z2, 1, 1);

            npk[z1][z2]--;
            npkSum[z1]--;
        }


        double[] pTable = new double[K1];

        for (int k = 0; k != K1; k++) {
            double expectTM = 1.0;
            int index = 0;
            for (int z2 : k2Count.keys()) {
                int x = a[k][z2] ? 1 : 0;
                int c = k2Count.get(z2);
                for (int i = 0; i != c; i++) {
                    //TODO:
                    expectTM *= (npk[k][z2] + x * alpha1 + alpha0 + i) / (K2 * alpha0 + a_sum[k] * alpha1 + npkSum[k] + index);
                    index++;
                }
            }

            pTable[k] = (mp[k] + alpha) / (M + K1 * alpha) * expectTM;
        }

        for (int k = 1; k != K1; k++) {
            pTable[k] += pTable[k - 1];
        }

        double r = rand.nextDouble() * pTable[K1 - 1];

        for (int k = 0; k != K1; k++) {
            if (pTable[k] > r) {
                z1 = k;
                break;
            }
        }

        mp[z1]++;
        for (int n = 0; n != N; n++) {
            int z2 = zAssigns_2[m][n];
            npk[z1][z2]++;
            npkSum[z1]++;
        }

        zAssigns_1[m] = z1;
    }

    public void sampleZ2(int m, int n) {
        int z1 = zAssigns_1[m];
        int z2 = zAssigns_2[m][n];
        int w = docs.get(m).get(n);

        npk[z1][z2]--;
        npkSum[z1]--;
        nkw[z2][w]--;
        nkwSum[z2]--;

        double VBeta = V * beta;

        double[] pTable = new double[K2];

        for (int k = 0; k != K2; k++) {
            //TODO:
            int x = a[z1][k] ? 1 : 0;
            pTable[k] = (npk[z1][k] + x * alpha1 + alpha0) / (npkSum[z1] + K2 * alpha0 + a_sum[z1] * alpha1) *
                    (nkw[k][w] + beta) / (nkwSum[k] + VBeta);
        }

        for (int k = 1; k != K2; k++) {
            pTable[k] += pTable[k - 1];
        }

        double r = rand.nextDouble() * pTable[K2 - 1];

        for (int k = 0; k != K2; k++) {
            if (pTable[k] > r) {
                z2 = k;
                break;
            }
        }

        npk[z1][z2]++;
        npkSum[z1]++;
        nkw[z2][w]++;
        nkwSum[z2]++;

        zAssigns_2[m][n] = z2;
        return;
    }

    public void sampleBinaryAMatrix() {
        int GIBBS_ITER = 40;

        a_sum = new int[K1];
        for (int m = 0; m != K1; m++) {
            for (int k = 0; k != K2; k++) {
                a[m][k] = npk[m][k] > 0;
                a_sum[m] += a[m][k] ? 1 : 0;
            }
        }

        double log_diff, ratio, p;
        for (int iter = 0; iter != GIBBS_ITER; iter++) {
            for (int m = 0; m != K1; m++) {
                for (int k = 0; k != K2; k++) {
                    if (a[m][k] && npk[m][k] == 0) {
                        log_diff = Gamma.logGamma(a_sum[m] * alpha1 + K2 * alpha0)
                                - Gamma.logGamma((a_sum[m] - 1) * alpha1 + K2 * alpha0);
                        log_diff -= Gamma.logGamma(npkSum[m] + a_sum[m] * alpha1 + K2 * alpha0)
                                - Gamma.logGamma(npkSum[m] + (a_sum[m] - 1) * alpha1 + K2 * alpha0);

                        ratio = Math.exp(log_diff) * pi_a[m] / (1.0 - pi_a[m]);
                        p = ratio / (1.0 + ratio);
                        if (rand.nextDouble() > p) {
                            a[m][k] = false;
                            a_sum[m]--;
                        }
                    } else if (!a[m][k]) {
                        log_diff = Gamma.logGamma((a_sum[m] + 1) * alpha1 + K2 * alpha0)
                                - Gamma.logGamma(a_sum[m] * alpha1 + K2 * alpha0);
                        log_diff -= Gamma.logGamma(npkSum[m] + (a_sum[m] + 1) * alpha1 + K2 * alpha0)
                                - Gamma.logGamma(npkSum[m] + a_sum[m] * alpha1 + K2 * alpha0);

                        ratio = Math.exp(log_diff) * pi_a[m] / (1.0 - pi_a[m]);
                        p = ratio / (1.0 + ratio);
                        if (rand.nextDouble() < p) {
                            a[m][k] = true;
                            a_sum[m]++;
                        }
                    }
                }

                BetaDistribution betaDist = new BetaDistribution(rand, gamma1 + a_sum[m], gamma0 + K2 - a_sum[m]);
                pi_a[m] = betaDist.sample();
            }
        }
    }

    public void sampleGamma1() {
        double shape = 1.0;
        double scale = 1.0;
        double sum_log_pi = 0.0;
        for (int k = 0; k != K1; k++) {
            sum_log_pi += Math.log(pi_a[k]);
        }
        double rate = 1.0 / scale - sum_log_pi;
        GammaDistribution gammaDist = new GammaDistribution(shape + K1, 1.0 / rate);
        gamma1 = gammaDist.sample();
    }

    private double log_factorial(int n, double a) {
        if (n == 0) return 0.0;
        double v = 0.0;
        for (int i = 0; i < n; i++) {
            v += Math.log(i + a);
        }
        return v;
    }

    public void sampleAlpha1() {
        int MH_STEP = 20;
        double s_var = 2.0;
        double shape = 1.0;
        double scale = 1.0;
        double u, tmp_alpha1, log_diff = 0.0;

        NormalDistribution gaussianDist = new NormalDistribution(0.0, s_var);

        for (int step = 0; step < MH_STEP; step++) {
            u = Math.log(rand.nextDouble());
            tmp_alpha1 = alpha1 + gaussianDist.sample();

            if (tmp_alpha1 > 0) {
                log_diff = (shape - 1) * (Math.log(tmp_alpha1 + alpha0) - Math.log(alpha1 + alpha0)) - (1.0 / scale) * (tmp_alpha1 - alpha1);
                for (int k1 = 0; k1 != K1; k1++) {
                    log_diff += Gamma.logGamma(tmp_alpha1 * a_sum[k1] + K2 * alpha0) - Gamma.logGamma(alpha1 * a_sum[k1] + K2 * alpha0);
                    log_diff -= Gamma.logGamma(tmp_alpha1 * a_sum[k1] + K2 * alpha0 + npkSum[k1])
                            - Gamma.logGamma(alpha1 * a_sum[k1] + K2 * alpha0 + npkSum[k1]);

                    for (int k2 = 0; k2 != K2; k2++) {
                        log_diff += log_factorial(npk[k1][k2], tmp_alpha1 + alpha0)
                                - log_factorial(npk[k1][k2], alpha1 + alpha0);
                    }
                }

                if (u < (log_diff < 0.0 ? log_diff : 0.0)) {
                    alpha1 = tmp_alpha1;
                }
            }
        }
    }

    public void estimate() {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("SPAM Iteration: " + iter + " ...");
            if (iter % this.saveStep == 0 && iter != 0 && iter != this.niters - 1) {
                this.storeResult(iter);
            }
            long begin = System.currentTimeMillis();
            for (int m = 0; m != M; m++) {
                this.sampleZ1(m);
            }
            for (int m = 0; m != M; m++) {
                int N = docs.get(m).size();
                for (int n = 0; n != N; n++) {
                    sampleZ2(m, n);
                }
            }

            if (iter % 1 == 0) {
                sampleBinaryAMatrix();
                System.out.print("\nPdoc sparsity : ");
                double sparsity = 0.0;
                for (int m = 0; m != K1; m++) {
                    sparsity += 1.0 - pi_a[m];
                }
                System.out.println((int) (100 * sparsity / K1));

                System.out.print("Pdoc real sparsity : ");
                sparsity = 0.0;
                for (int m = 0; m != K1; m++) {
                    sparsity += 1.0 - (double) a_sum[m] / K2;
                }
                System.out.println((int) (100 * sparsity / K1) + "\n");


                sampleGamma1();
                System.out.println("gamma1 : " + gamma1);


                sampleAlpha1();
                System.out.println("alpha1 : " + alpha1);
            }

//			if(iter!= 0 && iter % 500 == 0) {
//				this.saveTheta("C:/Users/zuoyuan/Desktop/Experiments/spam4st_news_theta_cw_"+iter+"_1.txt");
//			}

            System.out.println("Time(s) : " + (System.currentTimeMillis() - begin) / 1000.0 + "\n");
        }
        return;
    }

    public void printTopics(String path, int top_n) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(path);
        double[][] phi = computePhi();
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K2);
        for (int k = 0; k != K2; k++) {
            writer.write("Topic " + k + ":\n");
            for (int i = 0; i != top_n; i++) {
                writer.write(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue() + "\n");
            }
        }
        writer.close();
    }

    public double[][] computePhi() {
        double[][] phi = new double[K2][V];
        for (int k = 0; k != K2; k++) {
            for (int v = 0; v != V; v++) {
                phi[k][v] = (nkw[k][v] + beta) / (nkwSum[k] + V * beta);
            }
        }
        return phi;
    }

    int[][] ndk;
    int[] ndkSum;

    public void convert_zassigns_to_arrays_theta() {
        ndk = new int[M][K2];
        ndkSum = new int[M];

        for (int m = 0; m != M; m++) {
            for (int n = 0; n != docs.get(m).size(); n++) {
                ndk[m][zAssigns_2[m][n]]++;
                ndkSum[m]++;
            }
        }
    }

    public double[][] computeTheta() {
        convert_zassigns_to_arrays_theta();
        double[][] theta = new double[M][K2];
        for (int m = 0; m != M; m++) {
            int z1 = zAssigns_1[m];
            for (int k = 0; k != K2; k++) {
                int x = a[z1][k] ? 1 : 0;
                theta[m][k] = (ndk[m][k] + x * alpha1 + alpha0) / (ndkSum[m] + K2 * alpha0 + a_sum[z1] * alpha1);
            }
        }
        return theta;
    }

    public ArrayList<List<Entry<String, Double>>> sortedTopicWords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : w2i.keySet())
                term2weight.put(term, phi[k][w2i.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });
            res.add(pairs);
        }
        return res;
    }

    public void printTopics(int top_n) {
        double[][] phi = computePhi();
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K2);
        for (int k = 0; k != K2; k++) {
            System.out.println("Topic " + k + ":");
            for (int i = 0; i != top_n; i++) {
                System.out.println(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue());
            }
        }
    }

    public void savePhi(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] phi = computePhi();

        try {
            for (int k = 0; k != K2; k++) {
                for (int v = 0; v != V; v++) {
                    writer.append(phi[k][v] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveTheta(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] theta = computeTheta();

        try {
            for (int m = 0; m != M; m++) {
                for (int k = 0; k != K2; k++) {
                    writer.append(theta[m][k] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveWordmap(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (String word : w2i.keySet())
                writer.append(word + "\t" + w2i.get(word) + "\n");

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveAssign(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");
        try {
            for (int i = 0; i < zAssigns_2.length; i++) {
                for (int j = 0; j < zAssigns_2[i].length; j++) {
                    writer.write(docs.get(i).get(j) + ":" + zAssigns_2[i][j] + " ");
                }
                writer.write("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        return;
    }


    public void printModel() {
        System.out.println("K1 :" + this.K1 +
                "\tK2 :" + this.K2 +
                "\tniters :" + this.niters +
                "\tsaveStep :" + this.saveStep +
                "\talpha :" + this.alpha +
                "\talpha0 :" + this.alpha0 +
                "\talpha1 :" + this.alpha1 +
                "\tbeta :" + this.beta +
                "\tinputPath :" + this.inputPath +
                "\toutputPath :" + this.outputPath);
    }

    public void storeResult(int times) {
        String appendString = "final";
        if (times != 0) {
            appendString = times + "";
        }
        try {
            this.printTopics(outputPath + "/model-" + appendString + ".twords", 20);
            this.saveWordmap(outputPath + "/wordmap.txt");
            this.savePhi(outputPath + "/model-" + appendString + ".phi");
            this.saveTheta(outputPath + "/model-" + appendString + ".theta");
            this.saveAssign(outputPath + "/model-" + appendString + ".tassign");

        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    public void run() {
        printModel();
        this.loadTxts(inputPath);//
        this.initSTPAM();
        this.estimate();
        this.storeResult(0);

    }

    public static void SparsePAM4ST_V3(int P, int K, int iter, int saveStep, double alpha, double alpha0, double alpha1, double beta, int threadNum, String path) {
        ExecutorService service = Executors.newCachedThreadPool();
        File trainFile = new File(path);
        String parent_path = trainFile.getParentFile().getAbsolutePath();
        for (int i = 0; i < threadNum; i++) {
            (new File(parent_path + "/SparsePAM4ST_V3_" + P + "_" + K + "_" + alpha + "_" + alpha0 + "_" + alpha1 + "_" + beta + "/" + i)).mkdirs();
            service.submit(new SparsePAMForShortTexts_V3(P, K, iter, saveStep, alpha, alpha0, alpha1, beta, path, parent_path + "/SparsePAM4ST_V3_" + P + "_" + K + "_" + alpha + "_" + alpha0 + "_" + alpha1 + "_" + beta + "/" + i));

        }
        service.shutdown();
    }

//	public static void main(String args[]) {
//		
//		SparsePAMForShortTexts_V3 STPAM = new SparsePAMForShortTexts_V3();
//		STPAM.loadTxts("C:/Users/zuoyuan/Desktop/PAM4ST实验结果/MLDA/data2/news/en09062011.news.data.processed");//
//		STPAM.initSTPAM();
//		STPAM.estimate();
//		STPAM.saveTheta("C:/Users/zuoyuan/Desktop/Experiments/spam4st_news_theta_cw_final_1.txt");
//		STPAM.printTopics(10);
//		
//	}
}
